import re
from dataclasses import asdict, is_dataclass
from typing import Any, Dict, Optional, Union

import jinja2
from pydantic import BaseModel


class PromptTemplate:
    """
    提示模板类，用于管理和格式化提示文本。
    支持 JSON 和 Jinja2 两种模板格式。

    Attributes:
        template (str): 模板字符串
        format_type (str): 模板格式类型 ('json' 或 'jinja')
        variables (dict): 用于填充模板的变量
    """

    def __init__(self, template: str, format_type: str = 'json') -> None:
        """
        初始化提示模板。

        Args:
            template (str): 模板字符串
            format_type (str): 模板格式类型，默认为 'json'
        """
        self.template = template
        self.format_type = format_type

    def _convert_to_dict(
        self, variables: Optional[Union[Dict[str, str], BaseModel, Any]]
    ) -> Dict[str, str]:
        """
        将各种类型的变量转换为字典格式。

        Args:
            variables: 需要转换的变量，可以是字典、BaseModel 或数据类

        Returns:
            Dict[str, str]: 转换后的字典

        Raises:
            ValueError: 当变量类型不支持时抛出
        """
        if variables is None:
            return {}
        if isinstance(variables, BaseModel):
            return variables.dict()
        if is_dataclass(variables):
            return asdict(variables)
        if isinstance(variables, dict):
            return variables
        raise ValueError(
            'Unsupported variables type. Must be a dict, BaseModel, or '
            'dataclass.')

    def parse_template(self, template: str) -> Dict[str, str]:
        """
        从模板中提取变量。

        Args:
            template (str): 模板字符串

        Returns:
            Dict[str, str]: 包含所有变量名的字典，值均为 None
        """
        # 根据不同的格式类型使用不同的正则表达式提取变量
        if self.format_type == 'jinja':
            variables = re.findall(r'\{\{(.*?)\}\}', template)  # 匹配 Jinja2 格式 {{变量}}

        elif self.format_type == 'json':
            variables = re.findall(r'\{(.*?)\}', template)  # 匹配 JSON 格式 {变量}
            variables = [var for var in variables if '{' not in var]  # 过滤嵌套的花括号
        else:
            variables = []
        return {var.strip(): None for var in variables}

    def format_json(self, template: str, variables: Dict[str, str]) -> str:
        """
        使用 JSON 格式化模板。

        Args:
            template (str): JSON 格式的模板字符串
            variables (Dict[str, str]): 用于填充的变量

        Returns:
            str: 格式化后的字符串

        Raises:
            ValueError: 当模板格式无效时抛出
        """
        try:
            return template.format(**variables)
        except KeyError as e:
            raise ValueError('Invalid JSON template') from e

    def format_jinja(self, template: str, variables: Dict[str, str]) -> str:
        """
        使用 Jinja2 格式化模板。

        Args:
            template (str): Jinja2 格式的模板字符串
            variables (Dict[str, str]): 用于填充的变量

        Returns:
            str: 格式化后的字符串

        Raises:
            ValueError: 当模板格式无效时抛出
        """
        try:
            jinja_template = jinja2.Template(template)
            return jinja_template.render(variables)
        except jinja2.TemplateError as e:
            raise ValueError('Invalid Jinja template') from e

    def _update_variables_with_info(self) -> Dict[str, str]:
        """
        使用动作信息和代理信息更新变量字典。

        Returns:
            Dict[str, str]: 更新后的变量字典
        """
        variables = self.variables.copy()
        # 如果变量中没有但实例中有相关信息，则添加
        if 'action_info' not in variables and self.actions_info:
            variables['action_info'] = self.actions_info
        if 'agents_info' not in variables and self.agents_info:
            variables['agents_info'] = self.agents_info
        return variables

    def _check_variables_match(self, parsed_variables: Dict[str, str],
                               variables: Dict[str, str]) -> None:
        """
        检查所有变量是否都存在于解析出的变量中。

        Args:
            parsed_variables: 从模板解析出的变量
            variables: 要检查的变量

        Raises:
            ValueError: 当变量不匹配时抛出
        """
        if not all(key in parsed_variables for key in variables.keys()):
            raise ValueError(
                'Variables keys do not match the template variables')

    def format(
        self,
        **kwargs: Optional[Union[Dict[str, str], BaseModel, Any]],
    ) -> Any:
        """
        格式化模板的便捷方法。

        Args:
            **kwargs: 用于填充模板的关键字参数

        Returns:
            Any: 格式化后的字符串
        """
        self.variables = kwargs
        return str(self)

    def __str__(self) -> Any:
        """
        将模板转换为字符串，根据 format_type 选择相应的格式化方法。

        Returns:
            Any: 格式化后的字符串

        Raises:
            ValueError: 当格式类型不支持时抛出
        """
        parsed_variables = self.parse_template(self.template)
        updated_variables = self._update_variables_with_info()
        self._check_variables_match(parsed_variables, updated_variables)

        if self.format_type == 'json':
            return self.format_json(self.template, updated_variables)
        elif self.format_type == 'jinja':
            return self.format_jinja(self.template, updated_variables)
        else:
            raise ValueError('Unsupported format type')

    @property
    def actions_info(self) -> Optional[Dict[str, Any]]:
        """获取动作信息"""
        return getattr(self, '_action_info', None)

    @actions_info.setter
    def actions_info(self, value: Dict[str, Any]) -> None:
        """设置动作信息"""
        self._action_info = value

    @property
    def agents_info(self) -> Optional[Dict[str, Any]]:
        """获取代理信息"""
        return getattr(self, '_agents_info', None)

    @agents_info.setter
    def agents_info(self, value: Dict[str, Any]) -> None:
        """设置代理信息"""
        self._agents_info = value
